{-# LANGUAGE RecordWildCards #-}

module Te2.Te2_1 where

import Data.Bits (shiftL, testBit, (.|.))
import Data.Int (Int32)
import Text.Printf (printf)

------------------------------------------------------------
-- 模拟参数
------------------------------------------------------------
warpSize :: Int
warpSize = 8 -- 真实 GPU 里通常是 32，这里取 8 方便打印

type Mask = Int32 -- 位掩码，第 t 位为 1 表示线程 t 参与

-- 一条 warp 的所有线程数据
type WarpVec a = [a] -- 长度必须为 warpSize

------------------------------------------------------------
-- 纯函数版 shfl_up_sync
--   参数说明（和 CUDA 对应）：
--     mask   : 参与通信的线程掩码
--     delta  : 源线程 = 当前线程 - delta
--     val    : 当前线程的输入值
--     vec    : 整条 warp 的输入向量（用于读取源线程的值）
------------------------------------------------------------
shflUpSync :: (Eq a) => Mask -> Int -> a -> WarpVec a -> a
shflUpSync mask delta val vec
  | delta <= 0 = val -- delta<=0 无意义，CUDA 也返回原值
  | src < 0 || src >= warpSize = val -- 越界
  | not (mask `testBit` src) = val -- 源线程不在掩码里
  | otherwise = vec !! src
  where
    src = _threadIdx - delta
    _threadIdx = length (takeWhile (/= val) vec) -- 小技巧：用值反查下标
    -- 注意：这里为了演示方便，用“值”反查下标；真实场景请直接传下标。

-- 纯函数版 shfl_down_sync
shflDownSync :: (Eq a) => Mask -> Int -> a -> WarpVec a -> a
shflDownSync mask delta val vec
  | delta <= 0 = val
  | src < 0 || src >= warpSize = val
  | not (mask `testBit` src) = val
  | otherwise = vec !! src
  where
    src = _threadIdx + delta
    _threadIdx = length (takeWhile (/= val) vec)

------------------------------------------------------------
-- 工具：打印一条 warp 的结果
------------------------------------------------------------
printWarp :: (Show a) => String -> WarpVec a -> IO ()
printWarp title xs =
  putStrLn $ title ++ " = " ++ unwords (map (printf "%4s" . show) xs)

------------------------------------------------------------
-- 演示场景
------------------------------------------------------------
tt :: IO ()
tt = do
  let -- 初始数据：8 个线程分别握有自己的线程号
      inputData :: WarpVec Int
      inputData = [0 .. 7]

      -- 全部线程都参与
      fullMask :: Mask
      fullMask = foldr (.|.) 0 [1 `shiftL` t | t <- [0 .. warpSize - 1]]

      -- 1. shfl_up_sync 示例：delta=1，相当于每个线程拿“左边邻居”的值
      up1 = [shflUpSync fullMask 1 v inputData | v <- inputData]

      -- 2. shfl_down_sync 示例：delta=2，每个线程拿“右边隔一个”的值
      down2 = [shflDownSync fullMask 2 v inputData | v <- inputData]

      -- 3. 掩码示例：只让前 4 个线程参与，后 4 个线程虽然执行但得不到新数据
      low4Mask :: Mask
      low4Mask = foldr (.|.) 0 [1 `shiftL` t | t <- [0 .. 3]]
      up1Low4 = [shflUpSync low4Mask 1 v inputData | v <- inputData]

  printWarp "input" inputData
  printWarp "shfl_up 1" up1
  printWarp "shfl_down 2" down2
  printWarp "shfl_up 1 (low4 mask)" up1Low4

{-
\$ runhaskell ShflDemo.hs
input =    0    1    2    3    4    5    6    7
shfl_up 1 =    0    0    1    2    3    4    5    6
shfl_down 2 =    2    3    4    5    6    7    7    7
shfl_up 1 (low4 mask) =    0    0    1    2    4    5    6    7
-}